# 语义检索模块，目前只更改了test_system里面的
# 采用paraphrase-multilingual-MiniLM-L12-v2

import json
import os
import re
from typing import List, Dict
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
import logging
from performance_timer import global_timer, StepTimer

# 配置日志（默认设置为WARNING级别以减少性能影响）
logging.basicConfig(level=logging.WARNING, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 尝试导入jieba进行中文分词
try:
    import jieba
    import logging as jieba_logging
    
    JIEBA_AVAILABLE = True
    # 禁用jieba的日志输出以提高性能
    jieba_logging.basicConfig(level=logging.ERROR)
except ImportError:
    print("警告: jieba库未安装，将使用简单分词方法，对中文支持有限")
    JIEBA_AVAILABLE = False


# 下载faiss，numpy，sentence_transformers：pip install faiss-cpu numpy sentence-transformers

class SemanticRetriever:
    def __init__(self, kb_path='rust_knowledge_base/rust_docs_sample.json'):
        self.kb_path = kb_path
        self.knowledge_base = []
        self.index = None
        self.model = None
        
        # 添加缓存机制
        self.token_cache = {}
        self.embedding_cache = {}
        
        # 延迟加载标志
        self._kb_loaded = False
        self._model_loaded = False
        self._index_built = False
    
    def _ensure_resources_loaded(self):
        """
        确保所有必要的资源都已加载
        """
        if not self._kb_loaded:
            self.knowledge_base = self._load_knowledge_base()
            self._kb_loaded = True
        
        if not self._model_loaded:
            self._init_model()
            self._model_loaded = True
        
        if not self._index_built and self._kb_loaded and self._model_loaded:
            self._build_faiss_index()
            self._index_built = True
    
    def _init_model(self):
        """
        初始化嵌入模型
        """
        try:
            # 初始化嵌入模型
            # 优先尝试使用本地D盘的多语言模型
            local_model_path = 'D:/paraphrase-multilingual-MiniLM-L12-v2'
            if os.path.exists(local_model_path):
                print(f"使用本地多语言模型: {local_model_path}")
                self.model = SentenceTransformer(local_model_path)
            else:
                # 如果本地模型不存在，使用模型名称让系统自动下载并缓存
                print("使用模型名称，系统将自动查找缓存或下载多语言模型")
                self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
            
            # 禁用jieba的DEBUG输出
            if JIEBA_AVAILABLE:
                import jieba
                jieba.setLogLevel(20)  # 设置为INFO级别，隐藏DEBUG输出
        except Exception as e:
            logger.error(f"初始化嵌入模型失败: {str(e)}")
            print(f"警告: 初始化嵌入模型失败: {str(e)}")
            self.model = None

    def _load_knowledge_base(self) -> List[Dict]:  # 知识库加载并标准化
        """加载知识库并转换为统一格式"""
        if os.path.exists(self.kb_path):
            with open(self.kb_path, 'r', encoding='utf-8') as f:
                raw_docs = json.load(f)

            # 转换格式以匹配代码期望的结构
            converted_docs = []
            for i, doc in enumerate(raw_docs):
                # 优先使用知识库中已有的tags：如果存在且非空，直接使用；否则自动提取
                if 'tags' in doc and doc['tags']:  # 检查tags字段存在且不为空（空列表/空字符串都不满足）
                    tags = doc['tags']
                    # 可选：对已有tags做简单清洗（去重、过滤短词，保持数据一致性）
                    tags = list(dict.fromkeys([t.strip() for t in tags if len(t.strip()) > 1]))[:10]
                else:
                    # 无已有tags，自动从内容提取关键词作为标签
                    tags = []
                    if 'content' in doc:
                        # 简单的关键词提取，实际项目中可以使用更复杂的NLP方法
                        content_words = doc['content'].lower().split()
                        # 过滤掉常见停用词和短词
                        common_stopwords = set(
                            ['的', '了', '和', '是', '在', '有', '我', '他', '她', '它', '这', '那', '为', '与', '而',
                             '就', '都', '但', '及', '于', '要', '把', '将', '被', '也', '更', '还', '你', '您', '们'])
                        keywords = [word.strip('.,;!?\n') for word in content_words if
                                    word not in common_stopwords and len(word) > 2]
                        # 去重并取前10个关键词作为标签
                        tags = list(dict.fromkeys(keywords))[:10]

                # 优先使用知识库中已有的category：如果存在且非"unknown"，直接使用；否则自动分类
                category = doc.get('category', 'unknown')  # 先读取已有category，默认unknown
                if category == 'unknown' or not category:  # 仅当category为unknown或空时，执行自动分类
                    title_lower = doc.get('title', '').lower()
                    content_lower = doc.get('content', '').lower()

                    # 简单的分类规则
                    if any(word in title_lower or word in content_lower for word in
                           ['什么是', '什么', '定义', '概念', '介绍']):
                        category = 'definition'
                    elif any(word in title_lower or word in content_lower for word in
                             ['如何', '使用', '用法', '怎么', '步骤']):
                        category = 'usage'
                    elif any(word in title_lower or word in content_lower for word in
                             ['错误', '问题', 'debug', '修复', '解决']):
                        category = 'error_debug'
                    elif any(word in title_lower or word in content_lower for word in
                             ['比较', '区别', 'vs', 'vs.', '对比']):
                        category = 'comparison'
                    else:
                        category = 'faq'  # 默认归为常见问题

                # 标准化文档格式（统一字段）
                converted_doc = {
                    "id": f"k{i + 1:03d}",
                    "topic": doc.get("title", ""),  # 使用title作为topic
                    "content": doc.get("content", ""),
                    "tags": tags,  # 优先用已有tags，无则自动提取
                    "category": category,  # 优先用已有category，无则自动分类
                    "code": doc.get("code", "")
                }
                converted_docs.append(converted_doc)

            print(
                f"知识库加载完成，共{len(converted_docs)}条记录，平均每文档标签数: {sum(len(doc['tags']) for doc in converted_docs) / len(converted_docs):.1f}")
            return converted_docs
        else:
            # 若知识库文件不存在，返回默认知识库
            return [
                {
                    "id": "k001",
                    "topic": "生命周期",
                    "content": "Rust的生命周期用于确保引用始终有效。生命周期是引用保持有效的作用域。",
                    "tags": ["生命周期", "引用", "作用域"],
                    "category": "definition",
                    "code": "fn longest<'a>(x: &'a str, y: &'a str) -> &'a str {\n    if x.len() > y.len() { x } else { y }\n}"
                },
                {
                    "id": "k002",
                    "topic": "所有权",
                    "content": "Rust的所有权系统是内存安全的核心。每个值都有一个所有者，当所有者离开作用域时，值会被丢弃。",
                    "tags": ["所有权", "内存管理", "作用域"],
                    "category": "definition",
                    "code": "let s1 = String::from(\"hello\");\nlet s2 = s1; // s1的所有权移动到s2"
                },
                {
                    "id": "k003",
                    "topic": "可变变量",
                    "content": "使用mut关键字声明可变变量，允许修改变量的值。",
                    "tags": ["变量", "mut", "可变性"],
                    "category": "usage",
                    "code": "let mut x = 5;\nx = 10; // 可以修改"
                },
                {
                    "id": "k004",
                    "topic": "借用检查器错误",
                    "content": "借用检查器错误通常是由于违反了Rust的借用规则导致的。常见错误包括同时存在可变和不可变借用。",
                    "tags": ["借用检查器", "错误", "调试"],
                    "category": "error_debug",
                    "code": "let mut v = vec![1, 2, 3];\nlet first = &v[0];\nv.push(4); // 错误：同时存在可变和不可变借用"
                },
                {
                    "id": "k005",
                    "topic": "Rust vs C++",
                    "content": "Rust和C++都是系统编程语言，但Rust通过所有权系统提供内存安全，而C++需要手动管理内存。",
                    "tags": ["Rust", "C++", "比较"],
                    "category": "comparison",
                    "code": "// Rust: 自动内存管理\nlet s = String::from(\"hello\");\n// C++: 手动内存管理\n// std::string s = \"hello\";"
                }
            ]

    def _build_faiss_index(self):  # 构建 FAISS 向量索引
        """
        构建FAISS向量索引，使用更高效的索引类型，并支持索引持久化
        """
        if not self.knowledge_base or not self.model: # 若知识库或模型为空，直接返回
            return
        
        # 索引文件路径
        index_path = os.path.splitext(self.kb_path)[0] + '_index.faiss'
        
        # 尝试从文件加载索引，避免重复构建
        if os.path.exists(index_path) and not self._index_built:
            try:
                self.index = faiss.read_index(index_path)
                print(f"从文件加载FAISS索引: {index_path}")
                self._index_built = True
                return
            except Exception as e:
                print(f"加载FAISS索引失败，将重新构建: {str(e)}")

        # 使用缓存机制生成嵌入向量
        embeddings = []
        for doc in self.knowledge_base:
            # 组合关键信息：主题 + 分类 + 标签 + 内容 + 代码
            # 将标签转换为空格分隔的字符串
            tags_str = " ".join(doc.get('tags', [])) # 标签转为空格分隔的字符串

            # 构建完整文本（包含所有关键字段）
            combined_text = (
                f"主题: {doc['topic']} "
                f"分类: {doc['category']} "
                f"标签: {tags_str} "
                f"内容: {doc['content']}"
            )

            # 如果有代码也加入
            if doc.get('code'):
                combined_text += f" 代码: {doc['code']}"
            
            # 使用缓存的嵌入向量
            embedding = self._get_embedding(combined_text)
            if embedding is not None:
                embeddings.append(embedding)
        
        embeddings = np.array(embeddings).astype(np.float32)
        print(f"生成嵌入向量类型: {type(embeddings)}, 形状: {embeddings.shape}")

        # 根据文档数量选择合适的索引类型
        dimension = embeddings.shape[1]
        total_docs = len(embeddings)
        
        if total_docs > 1000:
            # 对于大数据集，使用IVF索引加速搜索
            nlist = min(100, total_docs // 10)
            quantizer = faiss.IndexFlatIP(dimension)
            self.index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT)
            self.index.train(embeddings)
        else:
            # 对于小数据集，使用Flat索引保证准确性
            self.index = faiss.IndexFlatIP(dimension)

        # 添加向量到索引
        self.index.add(embeddings)
        
        # 尝试保存索引到文件
        try:
            faiss.write_index(self.index, index_path)
            print(f"FAISS索引已保存到文件: {index_path}")
        except Exception as e:
            logger.warning(f"保存FAISS索引失败: {str(e)}")
            print(f"警告: 保存FAISS索引失败: {str(e)}")
            
        print(f"FAISS索引构建完成，共 {total_docs} 个文档")

    def _tokenize_text(self, text: str) -> List[str]:  # 文本分词
        """
        对文本进行分词，支持中英文，并使用缓存避免重复计算
        """
        # 检查缓存
        if text in self.token_cache:
            return self.token_cache[text]
        
        # 去除空白字符和特殊字符
        text = re.sub(r'\s+|[,.!?;\n]', ' ', text.lower().strip())

        # 检查文本中是否包含中文字符
        if JIEBA_AVAILABLE and re.search(r'[\u4e00-\u9fff]', text):
            # 使用jieba分词处理中文
            result = list(jieba.cut(text))
        else:
            # 英文分词（按空格）+ 简单中文分词（按字）
            # 提取所有汉字作为单独的词
            chinese_chars = re.findall(r'[\u4e00-\u9fff]', text)
            # 提取英文单词
            english_words = re.findall(r'[a-zA-Z]+', text)
            # 合并结果
            result = chinese_chars + english_words
        
        # 存入缓存
        self.token_cache[text] = result
        return result
    
    def _get_embedding(self, text: str):
        """
        获取文本的嵌入向量，并使用缓存避免重复计算
        """
        # 检查缓存
        if text in self.embedding_cache:
            return self.embedding_cache[text]
        
        # 确保模型已加载
        if not self._model_loaded:
            self._init_model()
            self._model_loaded = True
        
        if not self.model:
            return None
        
        # 生成嵌入向量
        embedding = self.model.encode([text], normalize_embeddings=True)[0]
        
        # 存入缓存
        self.embedding_cache[text] = embedding
        return embedding

    def _calculate_comprehensive_score(self, vector_score: float, doc: Dict, question: str, question_type: str, question_tokens: set = None) -> float:  #综合评分计算
        """
        计算综合评分，结合向量相似度和关键词匹配
        """
        # 向量相似度分数（已在向量检索时计算）
        vector_score_contribution = vector_score * 0.5  # 向量相似度权重0.5

        # 关键词匹配分数计算
        # 1. 提取文档的关键信息：主题、分类、标签、内容、代码
        doc_keywords = set()

        # 提取主题中的关键词
        if doc.get('topic'):
            doc_keywords.update(self._tokenize_text(doc['topic']))

        # 提取分类中的关键词
        if doc.get('category'):
            doc_keywords.update(self._tokenize_text(doc['category']))

        # 提取标签作为关键词
        if doc.get('tags'):
            for tag in doc['tags']:
                doc_keywords.update(self._tokenize_text(tag))

        # 提取内容中的关键词（取前100个字符进行分词）
        if doc.get('content'):
            content_preview = doc['content'][:100]
            doc_keywords.update(self._tokenize_text(content_preview))

        # 提取代码中的关键词（如果有）
        if doc.get('code'):
            # 从代码中提取标识符（变量名、函数名等）
            code_identifiers = set(re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', doc['code']))
            doc_keywords.update(code_identifiers)

        # 对问题进行分词（使用预计算的分词结果，如果提供）
        if question_tokens is None:
            question_keywords = set(self._tokenize_text(question))
        else:
            question_keywords = question_tokens

        # 计算关键词匹配比例
        keyword_match_score = 0.0
        if question_keywords:
            # 计算交集的大小（共同关键词的数量）
            common_keywords = question_keywords.intersection(doc_keywords)
            keyword_match_score = len(common_keywords) / len(question_keywords)

        # 关键词匹配分数贡献（权重0.5）
        keyword_match_contribution = keyword_match_score * 0.5

        # 类型匹配加分（可选，根据question_type）
        type_match_bonus = 0.0
        if question_type and doc.get('category') == question_type:
            type_match_bonus = 0.1  # 类型匹配加0.1分

        # 综合得分
        score = vector_score_contribution + keyword_match_contribution + type_match_bonus

        return min(1.0, score)  # 确保不超过1.0

    def retrieve(self, question: str, question_type: str = None, top_k: int = 3, debug: bool = False) -> List[Dict]:
        """
        语义检索相关文档 - 使用向量检索
        """
        try:
            with StepTimer("语义检索完整流程"):
                # 确保所有资源已加载
                with StepTimer("加载资源"):
                    self._ensure_resources_loaded()
                
                if not self.index:
                    print("FAISS索引未初始化")
                    return []

                # 使用缓存获取查询向量
                with StepTimer("生成查询向量"):
                    query_embedding = self._get_embedding(question)
                    if query_embedding is None:
                        return []
                    query_embedding = np.array([query_embedding]).astype(np.float32)

                # 动态计算搜索范围 - 优化搜索范围以提高性能
                total_docs = len(self.knowledge_base)
                # 减少搜索范围，从0.1减少到0.05，提高性能
                search_k = max(top_k * 5, min(int(total_docs * 0.05), total_docs))
                
                with StepTimer("FAISS检索"):
                    scores, indices = self.index.search(query_embedding, search_k)

                if debug:
                    print(f"检索问题: '{question}'")
                    print(f"搜索到 {len(indices[0])} 个候选结果")

                with StepTimer("结果过滤与排序"):
                    results = []
                    # 提前对问题进行分词（使用jieba，若可用），供后续标题匹配使用
                    question_tokens = set(self._tokenize_text(question))

                    # 限制处理的候选文档数量，避免不必要的计算
                    max_candidates = min(search_k, 20)
                    for i, (score, idx) in enumerate(zip(scores[0][:max_candidates], indices[0][:max_candidates])):
                        if idx < len(self.knowledge_base):
                            doc = self.knowledge_base[idx]

                            # 计算综合评分，传入预计算的问题分词结果以避免重复计算
                            final_score = self._calculate_comprehensive_score(score, doc, question, question_type, question_tokens)

                            # 核心：增加基于jieba分词的标题匹配逻辑
                            # 即使综合评分略低，但标题与问题高度相关，也纳入结果
                            is_relevant = False
                            # 原过滤条件：分数>0.35
                            if final_score > 0.35:
                                is_relevant = True
                            # 放宽条件：分数>0.25且标题与问题高度相关（依赖jieba分词结果）
                            elif final_score > 0.25:
                                # 对文档标题进行分词（使用jieba，若可用）
                                topic_tokens = set(self._tokenize_text(doc.get('topic', '')))
                                # 计算标题分词与问题分词的重叠比例，超过50%则认为高度相关
                                if topic_tokens and len(question_tokens & topic_tokens) / len(topic_tokens) > 0.5:
                                    if debug:
                                        print(f"DEBUG: 标题高度相关（分词匹配），破格录取: {doc.get('topic', '无标题')}")
                                    is_relevant = True

                            # 类型匹配条件（保持不变）
                            category_match = (not question_type or
                                              doc.get('category') == question_type or
                                              question_type == "unknown")

                            # 满足相关性和类型匹配则加入结果
                            if is_relevant and category_match:
                                results.append({
                                    **doc,
                                    'score': float(final_score)
                                })

                            # 提前结束如果已收集足够文档
                            if len(results) >= top_k:
                                break

                    # 按分数排序
                    results.sort(key=lambda x: x['score'], reverse=True)
                    
                global_timer.add_info("检索到文档数", len(results))
                if results:
                    global_timer.add_info("最高相似度得分", max(results, key=lambda x: x['score'])['score'])
                
                if debug:
                    print(f"返回 {len(results)} 个最终结果")
                return results
        except Exception as e:
            print(f"检索失败: {e}")
            global_timer.add_info("检索错误", str(e)[:100])
            return []

        # """语义检索相关文档"""
        # # 向量检索模式（优先）
        # if self.model and self.index:
        #     try:
        #         return self._vector_retrieve(question, question_type, top_k)
        #     except Exception as e:
        #         print(f"向量检索失败: {str(e)}")
        #         print("切换到备用检索方法: 关键词匹配")
        # else:
        #     print("模型或索引未初始化，使用备用检索方法: 关键词匹配")
        # # 备用关键词匹配模式
        # return self._keyword_retrieve(question, question_type, top_k)

    # def _vector_retrieve(self, question: str, question_type: str = None, top_k: int = 3) -> List[Dict]:
    #     """使用向量检索相关文档"""
    #     # 将问题转换为向量
    #     query_embedding = self.model.encode([question], normalize_embeddings=True).astype(np.float32)
    #
    #     # 使用FAISS进行相似度搜索 - 增加搜索范围
    #     # search_k = min(top_k * 3, len(self.knowledge_base))
    #     total_docs = len(self.knowledge_base)
    #     search_k = min(total_docs, top_k * 10)
    #     scores, indices = self.index.search(query_embedding, search_k)
    #
    #     print(f"检索问题: '{question}'")
    #     print(f"搜索到 {len(indices[0])} 个候选结果")
    #
    #     results = []
    #     for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
    #         if idx < len(self.knowledge_base):
    #             doc = self.knowledge_base[idx]
    #
    #             # 计算综合评分
    #             final_score = self._calculate_comprehensive_score(score, doc, question, question_type)
    #
    #             # 降低阈值并优化过滤逻辑
    #             # 1. 首先降低阈值到0.2，让更多文档能被检索到
    #             # 2. 对于特定情况可以进一步放宽限制
    #             is_relevant = False
    #
    #             # 基本条件：分数超过阈值
    #             if final_score > 0.2:
    #                 is_relevant = True
    #             # 如果问题和文档标题有较多重叠，即使分数略低也考虑
    #             elif final_score > 0.15:
    #                 question_tokens = set(self._tokenize_text(question))
    #                 topic_tokens = set(self._tokenize_text(doc.get('topic', '')))
    #                 if topic_tokens and len(question_tokens & topic_tokens) / len(topic_tokens) > 0.5:
    #                     print(f"DEBUG: 标题高度相关，破格录取: {doc.get('topic', '无标题')}")
    #                     is_relevant = True
    #
    #             # 分类匹配可以适当放宽
    #             category_match = not question_type or doc.get('category') == question_type or question_type == "unknown"
    #
    #             if is_relevant and category_match:
    #                 results.append({
    #                     **doc,
    #                     'score': float(final_score)
    #                 })
    #
    #             # 如果已经收集到足够的文档，提前结束
    #             if len(results) >= top_k:
    #                 break
    #
    #     # 按分数排序
    #     results.sort(key=lambda x: x['score'], reverse=True)
    #     print(f"返回 {len(results)} 个最终结果")
    #     return results
    #
    # def _keyword_retrieve(self, question: str, question_type: str = None, top_k: int = 3) -> List[Dict]:
    #     """备用关键词检索方法"""
    #     results = []
    #
    #     for doc in self.knowledge_base:
    #         score = 0
    #
    #         # 基于问题类型的过滤
    #         if question_type and doc.get('category') == question_type:
    #             score += 0.5
    #
    #         # 关键词匹配（使用改进的分词方法）
    #         question_tokens = set(self._tokenize_text(question))
    #         content_tokens = set(self._tokenize_text(doc['content']))
    #         tag_tokens = set()
    #         for tag in doc.get('tags', []):
    #             tag_tokens.update(self._tokenize_text(tag))
    #
    #         # 计算匹配度
    #         content_match = len(question_tokens & content_tokens) / len(question_tokens) if question_tokens else 0
    #         tag_match = len(question_tokens & tag_tokens) / len(question_tokens) if question_tokens else 0
    #
    #         score += content_match * 0.3 + tag_match * 0.2
    #
    #         if score > 0:
    #             results.append({
    #                 **doc,
    #                 'score': score
    #             })
    #
    #     # 按分数排序并返回top_k
    #     results.sort(key=lambda x: x['score'], reverse=True)
    #     return results[:top_k]

    def add_document(self, doc: Dict):
        """
        添加新文档到知识库
        """
        # 确保资源已加载
        self._ensure_resources_loaded()
        
        if 'id' not in doc:
            doc['id'] = f"k{len(self.knowledge_base) + 1:03d}"
        self.knowledge_base.append(doc)

        # 保存到文件
        try:
            with open(self.kb_path, 'w', encoding='utf-8') as f:
                json.dump(self.knowledge_base, f, ensure_ascii=False, indent=2)
            print(f"成功添加文档: {doc.get('topic', '无标题')}")
        except Exception as e:
            print(f"警告: 保存文档到文件失败: {str(e)}")

        # 重新构建索引（仅当向量检索模式可用时）
        try:
            if self.model:
                # 清除相关缓存以确保数据一致性
                self._index_built = False
                self.embedding_cache.clear()
                self._build_faiss_index()
        except Exception as e:
            print(f"警告: 重建FAISS索引失败: {str(e)}")

if __name__ == "__main__":
    print("开始测试语义检索模块...")
    # 创建检索器实例，默认不使用调试模式
    retriever = SemanticRetriever()

    # 测试问题
    test_questions = [
        "变量如何声明？",
        "什么是可变变量？",
        "如何定义函数？"
    ]

    # 启用调试模式进行测试
    for question in test_questions:
        print(f"\n=== 测试问题: {question} ===")
        results = retriever.retrieve(question, top_k=2, debug=True)

        for i, result in enumerate(results):
            print(f"{i + 1}. 主题: {result['topic']}, 分数: {result['score']:.4f}")
            print(f"   内容: {result['content']}")
            print()
